# -*- coding: UTF-8 -*-
"""
    @Author:YTQ
    @Time: 2022/7/20 16:40
    Description:datasets deal
    
"""

import os
import torch
from torchvision import transforms, datasets
import torch
import pandas as pd
import numpy as np
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
import torch.optim as optim
from PIL import Image


# 训练数据集加载 => 模型训练
class MyDataSet(Dataset):
    '''
        ogFolderPat：狗图片文件夹
        catFolderPath：猫图片文件夹
        data_transform：数据转换设置
        dogAct：狗标签数字
        catAct：猫标签数字
    '''
    def __init__(self, dogFolderPath, catFolderPath, data_transform, dogAct, catAct):
        self.imgPathArr = []
        self.labelArr = []
        # dog
        files = os.listdir(dogFolderPath)
        for f in files:
            self.imgPathArr.append(dogFolderPath+os.sep+f)
            self.labelArr.append(dogAct)
        # cat
        files.clear()
        files = os.listdir(catFolderPath)
        for f in files:
            self.imgPathArr.append(catFolderPath+os.sep+f)
            self.labelArr.append(catAct)

        self.transforms = data_transform

    # 图片及类别的获取
    def __getitem__(self, index):
        label = np.array(int(self.labelArr[index]))
        img_path = self.imgPathArr[index]
        pil_img = Image.open(img_path)
        if self.transforms:
            data = self.transforms(pil_img)
        else:
            pil_img = np.asarray(pil_img)
            data = torch.from_numpy(pil_img)
        return data, label

    def __len__(self):
        return len(self.imgPathArr)


# 验证数据集加载 => 模型验证
class TestDataSet(Dataset):
        '''
        filePath: 图片-类别对应文件（本例使用test.csv）
        image_path：用于验证的图片文件夹
        data_transform：数据转换设置
        dogAct：狗标签数字
        catAct：猫标签数字
    '''
    def __init__(self, filePath, image_path, data_transform, dogAct, catAct):
        self.imgPathArr = []
        self.labelArr = []
        # load csv file
        # 返回的是一个DataFrame数据
        pd_reader = pd.read_csv(filePath)
        length = len(pd_reader)
        for i in range(length):
            self.imgPathArr.append(f'{image_path}{os.sep}{pd_reader["id"][i]}.jpg')
            if pd_reader["label"][i] == "dog":
                self.labelArr.append(dogAct)
            elif pd_reader["label"][i] == "cat":
                self.labelArr.append(catAct)
        self.transforms = data_transform

    def __getitem__(self, index):
        label = np.array(int(self.labelArr[index]))
        img_path = self.imgPathArr[index]
        pil_img = Image.open(img_path)
        if self.transforms:
            data = self.transforms(pil_img)
        else:
            pil_img = np.asarray(pil_img)
            data = torch.from_numpy(pil_img)
        return data, label

    def __len__(self):
        return len(self.imgPathArr)


# 检测数据集加载 => 使用模型进行图片分类
class DetectDataSet(Dataset):
    '''
        filePath: 图片文件夹
        data_transform：数据转换设置
    '''
    def __init__(self, filePath, data_transform):
        self.imgPathArr = []
        self.labelArr = []
        if not os.path.exists(path=filePath):
            print("file not found")
        if os.path.isdir(filePath):
            for img in os.listdir():
                self.imgPathArr.append(filePath+os.sep+img)
                self.labelArr.append(-1)
        elif os.path.isfile(path=filePath):
            self.imgPathArr.append(filePath)
            self.labelArr.append(-1)
        self.transforms = data_transform

    def __getitem__(self, index):
        label = np.array(int(self.labelArr[index]))
        img_path = self.imgPathArr[index]
        pil_img = Image.open(img_path)
        if self.transforms:
            data = self.transforms(pil_img)
        else:
            pil_img = np.asarray(pil_img)
            data = torch.from_numpy(pil_img)
        return data, label

    def __len__(self):
        return len(self.imgPathArr)